import torch
import os
import atexit


class TensorboardWriter:
    def __init__(self, output_dir):
        from torch.utils.tensorboard import SummaryWriter

        log_dir = os.path.join(output_dir, 'logs')
        os.makedirs(log_dir, exist_ok=True)
        self.writer = SummaryWriter(log_dir=log_dir)
        atexit.register(self.close)

    def add_dict(self, metric_dict, step, prefix=''):
        for key, value in metric_dict.items():
            self.writer.add_scalar(prefix + key, value, global_step=step)

    def close(self):
        self.writer.close()

    def __del__(self):
        self.close()
